package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"os"
	"os/exec"
	"path"
	"slices"
	"strings"
	"text/template"

	"github.com/gomlx/gomlx/internal/must"
	"github.com/gomlx/gomlx/internal/xlabuilderparser"
	"github.com/gomlx/gomlx/pkg/support/sets"
	"github.com/gomlx/gomlx/pkg/support/xslices"
	"k8s.io/klog/v2"
)

const (
	standardOpsXlaFile = "gen_standard_ops.go"
)

var (
	// methodsToExclude from support. Typically, deprecated functions.
	methodsToExclude = sets.MakeWith(
		"And", "Or", "Xor", "Not", "ReduceAnd", "ReduceOr", "ReduceXor", "ScatterAdd",
		"SelectAndScatterSum", // SelectAndScatterSum is deprecated and doesn't work.
	)

	standardOpsTemplate = template.Must(template.New(standardOpsXlaFile).Parse(`
/***** File generated by ./internal/cmd/xla_generator, based on github.com/gomlx/gopjrt. Don't edit it directly. *****/

package xla

import (
	"github.com/gomlx/gomlx/backends"
	"github.com/gomlx/gomlx/pkg/core/shapes"
	"github.com/gomlx/gopjrt/dtypes"
	"github.com/gomlx/gopjrt/xlabuilder"
	"github.com/pkg/errors"
)

{{- range .}}
{{- range .Comments}}
// {{.}}
{{- end}}
func (b *Builder) {{.Name}}({{range .Parameters}}{{.Names}} {{.Type}}, {{end}}) (backends.Op, error) {
{{- range .ValuesSpec}}
{{.}}
{{- end}}
	xla_result, err := xlabuilder.{{.Name}}({{range.Values}}{{.}}, {{end}})
	if err != nil {
		return nil, errors.WithMessagef(err, "Backend %q: failed {{.Name}}", BackendName)
	}
	return xla_result, nil
}
{{end}}
`))

	convertOpTemplate = template.Must(template.New("convertOp").Parse(
		`	xla_{{.Param}}, err := b.verifyAndCastOp({{.Param}}, "{{.Param}}")
	if err != nil {
		return nil, errors.WithMessagef(err, "Backend %q: failed {{.Func}}", BackendName)
	}
`))

	convertOpListTemplate = template.Must(template.New("convertOpList").Parse(
		`	xla_{{.Param}}, err := b.verifyAndCastOps({{.Param}}, "{{.Param}}")
	if err != nil {
		return nil, errors.WithMessagef(err, "Backend %q: failed {{.Func}}", BackendName)
	}
`))

	convertShapeTemplate = template.Must(template.New("convertShape").Parse(
		`	xla_{{.Param}} := shapeToXShape({{.Param}})`))

	convertConvolveAxesConfigTemplate = template.Must(template.New("convertConvolveAxesConfig").Parse(
		`	xla_{{.Param}} := convertConvolveAxesConfig({{.Param}})`))

	convertPadAxisListTemplate = template.Must(template.New("convertPadAxisList").Parse(
		`	var xla_{{.Param}} []xlabuilder.PadAxis
	for _, pad := range {{.Param}} {
		xla_{{.Param}} = append(xla_{{.Param}}, convertPadAxis(pad))
	}`))

	convertFFTTypeTemplate = template.Must(template.New("convertFFTType").Parse(
		`	var xla_{{.Param}} = convertFFTType({{.Param}})`))
)

type FuncInfo struct {
	Name       string
	Parameters []*ParamInfo
	ValuesSpec []string
	Values     []string
	Comments   []string
}

type ParamInfo struct {
	Names, Type string
}

func selectAndScatterComments(funcName string) []string {
	aggregation := funcName[len(funcName)-3:]
	return []string{
		funcName + " runs windows (similar to ReduceWindow) over the operand, selects values to updates the output (like ScatterAdd)",
		"It selects the values in the window such that it works as reverse for Scatter" + aggregation + ".",
		"See details in https://openxla.org/xla/operation_semantics#selectandscatter",
	}
}

func execTemplate(tmpl *template.Template, data any) string {
	buf := &bytes.Buffer{}
	must.M(tmpl.Execute(buf, data))
	return buf.String()
}

// GenerateStandardOpsImplementation generates the XLA implementation for the various standard ops, that just calls
// the underlying xlabuilder function.
// The rest of the ops are maintained manually.
func GenerateStandardOpsImplementation(extractor *xlabuilderparser.NodeTextExtractor, xlaBuilderPkg *ast.Package) {
	var standardOps []FuncInfo
	xlabuilderparser.EnumerateStandardOpsFunctions(extractor, xlaBuilderPkg, func(funcDecl *ast.FuncDecl) {
		if methodsToExclude.Has(funcDecl.Name.Name) {
			return
		}
		fi := FuncInfo{
			Name: funcDecl.Name.Name,
			Comments: slices.DeleteFunc(
				strings.Split(funcDecl.Doc.Text(), "\n"),
				func(s string) bool {
					return s == "" || s == "The op is created on the same XlaBuilder as used for x." ||
						strings.Index(s, "ScatterCustom") != -1
				}),
		}
		for _, param := range funcDecl.Type.Params.List {
			if param.Names[0].Name == "builder" {
				// builder is not needed by the API, since the builder is given by the object implementing the interface.
				fi.Values = append(fi.Values, "b.builder")
				continue
			}
			names := xslices.Map(param.Names, func(ident *ast.Ident) string { return ident.Name })
			pi := &ParamInfo{
				Names: strings.Join(names, ", "),
				Type:  extractor.Get(param.Type),
			}
			if pi.Type == "*Op" {
				pi.Type = "backends.Op"
			} else if pi.Type == "...*Op" {
				pi.Type = "...backends.Op"
			} else if pi.Type == "[]*Op" {
				pi.Type = "[]backends.Op"
			} else if pi.Type == "Shape" {
				pi.Type = "shapes.Shape"
			} else if pi.Type == "ConvolveAxesConfig" {
				pi.Type = "backends.ConvolveAxesConfig"
			} else if pi.Type == "...PadAxis" {
				pi.Type = "...backends.PadAxis"
			} else if pi.Type == "xla_data.FftType" || pi.Type == "FFTType" {
				pi.Type = "backends.FFTType"
			}
			fi.Parameters = append(fi.Parameters, pi)
			valueSpecArgs := struct{ Param, Func string }{Func: fi.Name}

			// Encode conversions from backends.Builder API to xlabuilder API.
			if pi.Type == "backends.Op" {
				for ii, name := range names {
					valueSpecArgs.Param = name
					fi.ValuesSpec = append(fi.ValuesSpec, execTemplate(convertOpTemplate, valueSpecArgs))
					names[ii] = "xla_" + name
				}
			} else if pi.Type == "shapes.Shape" {
				for ii, name := range names {
					valueSpecArgs.Param = name
					fi.ValuesSpec = append(fi.ValuesSpec, execTemplate(convertShapeTemplate, valueSpecArgs))
					names[ii] = "xla_" + name
				}
			} else if pi.Type == "...backends.Op" {
				for ii, name := range names {
					valueSpecArgs.Param = name
					fi.ValuesSpec = append(fi.ValuesSpec, execTemplate(convertOpListTemplate, valueSpecArgs))
					names[ii] = "xla_" + name + "..."
				}
			} else if pi.Type == "[]backends.Op" {
				for ii, name := range names {
					valueSpecArgs.Param = name
					fi.ValuesSpec = append(fi.ValuesSpec, execTemplate(convertOpListTemplate, valueSpecArgs))
					names[ii] = "xla_" + name
				}
			} else if pi.Type == "backends.ConvolveAxesConfig" {
				for ii, name := range names {
					valueSpecArgs.Param = name
					fi.ValuesSpec = append(fi.ValuesSpec, execTemplate(convertConvolveAxesConfigTemplate, valueSpecArgs))
					names[ii] = "xla_" + name
				}
			} else if pi.Type == "...backends.PadAxis" {
				for ii, name := range names {
					valueSpecArgs.Param = name
					fi.ValuesSpec = append(fi.ValuesSpec, execTemplate(convertPadAxisListTemplate, valueSpecArgs))
					names[ii] = "xla_" + name + "..."
				}
			} else if pi.Type == "backends.FFTType" {
				for ii, name := range names {
					valueSpecArgs.Param = name
					fi.ValuesSpec = append(fi.ValuesSpec, execTemplate(convertFFTTypeTemplate, valueSpecArgs))
					names[ii] = "xla_" + name
				}
			} else if strings.HasPrefix(pi.Type, "...") {
				for ii, name := range names {
					names[ii] = name + "..."
				}
			}
			fi.Values = append(fi.Values, names...)
		}
		if strings.HasPrefix(fi.Name, "SelectAndScatter") {
			fi.Comments = selectAndScatterComments(fi.Name)
		}
		standardOps = append(standardOps, fi)
	})
	slices.SortFunc(standardOps, func(a, b FuncInfo) int { return strings.Compare(a.Name, b.Name) })

	fileName := path.Join(must.M1(os.Getwd()), standardOpsXlaFile)
	f := must.M1(os.Create(fileName))

	must.M(standardOpsTemplate.Execute(f, standardOps))
	cmd := exec.Command("go", "fmt", fileName)
	klog.V(1).Infof("\t%s\n", cmd)
	cmd.Stderr = os.Stderr
	must.M(cmd.Run())
	fmt.Printf("✅ xla_generator:         \tsuccessfully generated %s\n", fileName)
}
